package com.treasure.domain.iqap;

import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.treasure.common.model.ChatResponse;
import com.treasure.common.model.command.ConversationRequest;
import com.treasure.common.types.ClientId;
import com.treasure.common.types.InteractionId;
import com.treasure.domain.domain.ChatHistoryDomainService;
import com.treasure.domain.domain.LargeLanguageModelDomainService;
import com.treasure.domain.domain.RelatedTextExtractionDomainService;
import com.treasure.common.model.ChatHistoryEntity;
import com.treasure.domain.domain.TextVariablesCombinedDomainService;
import com.treasure.domain.domain.entity.ChatHistoryQuery;
import com.treasure.domain.domain.entity.PromptTemplateContext;
import com.treasure.domain.domain.entity.QuestionRequest;
import com.treasure.domain.domain.entity.SegmentDocEntity;
import com.treasure.common.types.Prompt;
import com.treasure.common.types.Question;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Service;

import java.io.IOException;
import java.util.Collections;
import java.util.List;
import java.util.Optional;

/**
 * @author sunhongmin
 * @date 2023/8/23 20:05
 * @description
 */
@RequiredArgsConstructor
@Service
@Slf4j
public class IntelligentAnswerProcessorImpl implements IntelligentAnswerProcessor {

    public static final int MEMORY_ROUNDS = 2;
    private final LargeLanguageModelDomainService largeLanguageModelDomainService;
    private final RelatedTextExtractionDomainService relatedTextExtractionDomainService;
    private final ChatHistoryDomainService chatHistoryDomainService;
    private final TextVariablesCombinedDomainService textVariablesCombinedDomainService;

    @Override
    public String questionAndAnswer(String question, String promptTemplate) throws IOException {
        // 关联匹配
        List<SegmentDocEntity> segmentDocList = relatedTextExtractionDomainService.relatedTextExtraction(new Question(question));
        log.info("用户输入: \n{}", question);
        // prompt整合
        String prompt = textVariablesCombinedDomainService.promptTemplate(new PromptTemplateContext(segmentDocList, question));

        log.info("整合后的prompt： \n {}", prompt);
        // 调用大语言模型
        String result = largeLanguageModelDomainService.callGeneration(new Prompt(prompt));
        log.info("回答： \n {}", result);
        return result;
    }

    @Override
    public ChatResponse conversationForChain(QuestionRequest questionRequest, Boolean syncCall) throws IOException, NoApiKeyException, InputRequiredException {
        // 获取会话ID
        Integer interactionId = getOrGenerateInteractionId(questionRequest);
        questionRequest.setInteractionId(interactionId);

        // 文档关联匹配
        List<SegmentDocEntity> segmentDocList = relatedTextExtractionDomainService.relatedTextExtraction(new Question(questionRequest.getQuestion()));
        log.info("用户输入: \n{}", questionRequest.getQuestion());

        // prompt整合
        String prompt = textVariablesCombinedDomainService.promptTemplate(new PromptTemplateContext(segmentDocList, questionRequest.getQuestion()));
//        log.info("整合后的prompt： \n {}", prompt);

        // 历史聊天嵌入
        ChatHistoryEntity chatHistoryList = chatHistoryDomainService.getChatHistoryList(new ChatHistoryQuery(new InteractionId(interactionId), MEMORY_ROUNDS));
        List<ChatHistoryEntity.ChatHistoryBean> chatHistoryBeans = Optional.ofNullable(chatHistoryList).map(ChatHistoryEntity::getChatHistoryBeans).orElse(Collections.emptyList());

        // 调用大语言模型
        ConversationRequest conversationRequest = new ConversationRequest();
        conversationRequest
                .setInteractionId(new InteractionId(questionRequest.getInteractionId()))
                .setQuestion(new Question(questionRequest.getQuestion()))
                .setPrompt(new Prompt(prompt))
                .setChatHistoryList(chatHistoryBeans)
                .setSseClientId(new ClientId(questionRequest.getClientToken()))
                .setSyncCall(syncCall);
        String resMessage = largeLanguageModelDomainService.callGeneration(conversationRequest);

        ChatResponse chatResponse = new ChatResponse();
        chatResponse.setInteractionId(interactionId);
        chatResponse.setResponseMessage(resMessage);
        return chatResponse;
    }

    private Integer getOrGenerateInteractionId(QuestionRequest questionRequest) {
        // 校验会话是否存在
        boolean interactionExists = chatHistoryDomainService.checkInteractionExists(questionRequest.getInteractionId());
        Integer interactionId;
        if (!interactionExists) {
            // 新建会话
            interactionId = chatHistoryDomainService.createInteractionId(new Question(questionRequest.getQuestion())).getValue();
        } else {
            interactionId = questionRequest.getInteractionId();
        }
        return interactionId;
    }

}
